"""
统一MCP服务器 - 单一端点提供多个工具

在单一端口（8000）提供所有MCP工具：
- Python代码执行工具
- MySQL数据库查询工具

客户端只需连接一个端点即可使用所有功能
"""

import multiprocessing
from typing import Any

import mcp.types as types
from mcp.server import Server
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.routing import Route, Mount
from starlette.requests import Request
from starlette.responses import JSONResponse
import uvicorn

# 导入MCP组件
from mcp_component import BaseMCPComponent, PythonExecutorComponent, MySQLQueryComponent


# ============================================================================
# 全局配置
# ============================================================================

# 组件实例
python_component = PythonExecutorComponent()
mysql_component = MySQLQueryComponent()

# 所有组件列表（统一管理）
ALL_COMPONENTS = [python_component, mysql_component]

# 工具名称到组件的映射缓存（O(1)查找）
TOOL_TO_COMPONENT_CACHE: dict[str, BaseMCPComponent] = {}


def build_tool_cache():
    """构建工具名称到组件的映射缓存"""
    TOOL_TO_COMPONENT_CACHE.clear()
    for component in ALL_COMPONENTS:
        if component.is_available():
            for tool in component.get_tools():
                TOOL_TO_COMPONENT_CACHE[tool.name] = component
build_tool_cache()

# MCP服务器实例
mcp_server = Server("unified-mcp-server")
sse_transport = SseServerTransport("/")  # Mount会截断/sse前缀，所以这里用/


# ============================================================================
# MCP工具注册
# ============================================================================

@mcp_server.list_tools()
async def list_tools() -> list[types.Tool]:
    """列出所有可用的工具"""    
    tools = []
    for component in ALL_COMPONENTS:
        if component.is_available():
            tools.extend(component.get_tools())
    
    return tools


@mcp_server.call_tool()
async def call_tool(name: str, arguments: Any) -> list[types.TextContent]:
    """处理工具调用 - O(1) 快速查找"""
    print(f"调用工具: {name}")
    
    # 从缓存中快速查找对应的组件
    component = TOOL_TO_COMPONENT_CACHE.get(name)
    
    if component:
        return await component.handle_tool_call(name, arguments)
    
    # 未找到工具
    return [types.TextContent(type="text", text=f"未知工具: {name}")]


# ============================================================================
# HTTP服务器
# ============================================================================

async def health_check(request: Request):
    """健康检查端点"""
    return JSONResponse({
        "status": "healthy",
        "service": "unified-mcp-server",
        "version": "2.0.0",
        "port": 8000,
        "features": {
            "python_executor": python_component.is_available(),
            "mysql_query": mysql_component.is_available(),
        },
        "tools_count": len(await list_tools())
    })

async def handle_mcp_endpoint(scope, receive, send):
    """统一的MCP端点 - Session由SseServerTransport内部管理"""
    method = scope["method"]
    
    if method == "GET":
        async with sse_transport.connect_sse(scope, receive, send) as streams:
            await mcp_server.run(
                streams[0],
                streams[1],
                mcp_server.create_initialization_options()
            )
    elif method == "POST":
        await sse_transport.handle_post_message(scope, receive, send)
    else:
        await send({
            "type": "http.response.start",
            "status": 405,
            "headers": [[b"content-type", b"text/plain"]],
        })
        await send({
            "type": "http.response.body",
            "body": b"Method Not Allowed",
        })


app = Starlette(
    debug=False,
    routes=[
        Route("/health", health_check, methods=["GET"]),
        Mount("/sse", app=handle_mcp_endpoint),  # ASGI 应用
    ]
)


def main():
    """主函数"""
    # Windows多进程支持
    multiprocessing.freeze_support()
    
    # 启动服务
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info",
        limit_concurrency=1000
    )

if __name__ == "__main__":
    main()
